import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

from ultralytics import YOLO
import cv2
import matplotlib.pyplot as plt
from matplotlib import rcParams
from collections import Counter

# 设置支持中文的字体
rcParams['font.family'] = 'sans-serif'
rcParams['font.sans-serif'] = ['SimHei']  # 或者使用其他支持中文的字体

def infer(frame, model, confidence_threshold=0.8):
    # 执行推理
    results = model(frame)

    # 统计每个类别的数量
    class_counts = Counter()

    # 输出检测的类别、置信度和边界框信息，过滤置信度低于阈值的检测框
    filtered_boxes = []
    for detection in results[0].boxes.data:
        x1, y1, x2, y2, score, class_id = detection.tolist()
        if score >= confidence_threshold:  # 只处理置信度高于设定阈值的框
            class_name = model.names[int(class_id)]
            class_counts[class_name] += 1
            filtered_boxes.append(detection)

    # 如果没有满足条件的检测结果，直接返回原始帧
    if not filtered_boxes:
        return frame

    # 可视化检测结果（过滤后的结果）
    results[0].boxes.data = filtered_boxes  # 更新检测结果为过滤后的框
    result_frame = results[0].plot()  # 绘制检测结果
    return result_frame

def main():
    # 加载训练好的YOLO模型
    model = YOLO('./runs/detect/yolo_custom3/weights/best.pt')  # 使用训练后保存的模型

    # 打开摄像头
    cap = cv2.VideoCapture(1)  # 参数 0 表示使用默认摄像头

    if not cap.isOpened():
        print("无法打开摄像头")
        return

    plt.ion()  # 开启交互模式

    while True:
        ret, frame = cap.read()  # 从摄像头读取一帧
        if not ret:
            print("无法读取摄像头画面")
            break

        # 将图像转换为RGB模式
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        # 推理并过滤置信度小于 0.8 的框
        result_frame = infer(frame_rgb, model, confidence_threshold=0.2)

        # 使用 matplotlib 显示图像
        plt.imshow(result_frame)
        plt.title('实时检测')
        plt.axis('off')
        plt.pause(0.001)  # 更新显示

        # 按下 'q' 键退出
        if plt.waitforbuttonpress(0.001):
            break

    # 释放摄像头并关闭所有窗口
    cap.release()
    plt.ioff()  # 关闭交互模式
    plt.close()

if __name__ == '__main__':
    main()
